"""
EmbedTrack 训练和推理速度优化配置
包含所有可以提高性能的设置和优化策略
"""

import os
import torch
import psutil
from typing import Dict, Any

class OptimizationConfig:
    """优化配置类"""
    
    def __init__(self):
        self.setup_environment()
        self.setup_cuda_optimizations()
        self.setup_memory_optimizations()
        
    def setup_environment(self):
        """设置环境变量优化"""
        # CPU线程优化
        cpu_count = psutil.cpu_count()
        optimal_threads = min(cpu_count // 2, 8)
        
        os.environ['OMP_NUM_THREADS'] = str(optimal_threads)
        os.environ['MKL_NUM_THREADS'] = str(optimal_threads)
        os.environ['NUMEXPR_NUM_THREADS'] = str(optimal_threads)
        os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
        
        # 禁用不必要的警告
        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
        
    def setup_cuda_optimizations(self):
        """设置CUDA优化"""
        if torch.cuda.is_available():
            # 清空GPU缓存
            torch.cuda.empty_cache()
            
            # cuDNN优化
            torch.backends.cudnn.benchmark = True
            torch.backends.cudnn.deterministic = False
            torch.backends.cudnn.enabled = True
            
            # 内存分配策略
            torch.cuda.set_per_process_memory_fraction(0.95)
            
            # 启用混合精度训练
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True
            
    def setup_memory_optimizations(self):
        """设置内存优化"""
        # 设置PyTorch内存分配器
        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
        
    def get_training_optimizations(self) -> Dict[str, Any]:
        """获取训练优化配置"""
        return {
            # 批次大小优化
            'batch_size_multiplier': 2 if torch.cuda.is_available() else 1,
            'virtual_batch_multiplier': 2,
            
            # 数据加载优化
            'num_workers': min(psutil.cpu_count() // 2, 8),
            'pin_memory': True,
            'persistent_workers': True,
            'prefetch_factor': 2,
            
            # 模型优化
            'use_amp': True,  # 自动混合精度
            'gradient_accumulation_steps': 2,
            'max_grad_norm': 1.0,
            
            # 学习率优化
            'learning_rate': 1e-3,
            'weight_decay': 1e-4,
            
            # 早停和检查点
            'early_stopping_patience': 10,
            'save_checkpoint_freq': 5,
            
            # 数据增强优化
            'use_fast_augmentation': True,
            'cache_augmented_data': True,
        }
    
    def get_inference_optimizations(self) -> Dict[str, Any]:
        """获取推理优化配置"""
        return {
            # 批次大小
            'batch_size': 32 if torch.cuda.is_available() else 8,
            
            # 数据加载
            'num_workers': min(psutil.cpu_count() // 2, 4),
            'pin_memory': True,
            
            # 模型优化
            'use_amp': True,
            'use_jit': True,  # 使用TorchScript
            'fuse_bn': True,  # 融合批归一化层
            
            # 推理优化
            'use_tensorrt': False,  # 如果有TensorRT可用
            'optimize_for_inference': True,
            
            # 内存优化
            'clear_cache_freq': 10,
            'use_gradient_checkpointing': False,  # 推理时不需要
        }
    
    def get_data_optimizations(self) -> Dict[str, Any]:
        """获取数据优化配置"""
        return {
            # 数据预处理
            'use_preprocessing_cache': True,
            'cache_dir': './data_cache',
            'preload_data': True,
            
            # 图像处理
            'use_fast_resize': True,
            'use_fast_normalization': True,
            
            # 数据增强
            'use_fast_augmentation': True,
            'augmentation_cache_size': 1000,
            
            # 文件I/O优化
            'use_memory_mapping': True,
            'buffer_size': 8192,
        }
    
    def get_model_optimizations(self) -> Dict[str, Any]:
        """获取模型优化配置"""
        return {
            # 模型结构优化
            'use_channels_last': True,  # 使用channels_last内存格式
            'use_optimized_conv': True,
            
            # 量化优化
            'use_quantization': False,  # 如果需要可以启用
            'quantization_type': 'int8',
            
            # 模型剪枝
            'use_pruning': False,
            'pruning_ratio': 0.3,
            
            # 知识蒸馏
            'use_knowledge_distillation': False,
        }

# 全局优化配置实例
optimization_config = OptimizationConfig()

def get_optimized_training_config():
    """获取优化的训练配置"""
    return optimization_config.get_training_optimizations()

def get_optimized_inference_config():
    """获取优化的推理配置"""
    return optimization_config.get_inference_optimizations()

def get_optimized_data_config():
    """获取优化的数据配置"""
    return optimization_config.get_data_optimizations()

def get_optimized_model_config():
    """获取优化的模型配置"""
    return optimization_config.get_model_optimizations()

# 预设配置
PRESET_CONFIGS = {
    'fast_training': {
        'batch_size': 8,
        'num_workers': 4,
        'use_amp': True,
        'gradient_accumulation_steps': 4,
        'learning_rate': 2e-3,
        'early_stopping_patience': 5,
    },
    'fast_inference': {
        'batch_size': 64,
        'num_workers': 2,
        'use_amp': True,
        'use_jit': True,
        'fuse_bn': True,
    },
    'memory_efficient': {
        'batch_size': 2,
        'num_workers': 2,
        'use_amp': True,
        'gradient_checkpointing': True,
        'clear_cache_freq': 5,
    },
    'high_accuracy': {
        'batch_size': 4,
        'num_workers': 6,
        'use_amp': True,
        'learning_rate': 5e-4,
        'early_stopping_patience': 15,
    }
}

def get_preset_config(preset_name: str) -> Dict[str, Any]:
    """获取预设配置"""
    return PRESET_CONFIGS.get(preset_name, PRESET_CONFIGS['fast_training']) 